/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package org.xhy.sniffer.datacarrier;


import org.xhy.sniffer.datacarrier.buffer.BufferStrategy;
import org.xhy.sniffer.datacarrier.buffer.Channels;
import org.xhy.sniffer.datacarrier.consumer.ConsumeDriver;
import org.xhy.sniffer.datacarrier.consumer.ConsumerPool;
import org.xhy.sniffer.datacarrier.consumer.IConsumer;
import org.xhy.sniffer.datacarrier.consumer.IDriver;
import org.xhy.sniffer.datacarrier.partition.IDataPartitioner;
import org.xhy.sniffer.datacarrier.partition.SimpleRollingPartitioner;

import java.util.Properties;

/**
 * DataCarrier 主类。使用此实例来设置生产者/消费者模型。
 */
public class DataCarrier<T> {
    private Channels<T> channels; // 保存数据的通道，支持多通道操作
    private IDriver driver; // 控制消费者运行的驱动
    private String name; // 当前实例的名称

    // 构造函数，创建具有默认名称的 DataCarrier 实例，指定通道和缓冲区大小
    public DataCarrier(int channelSize, int bufferSize) {
        this("DEFAULT", channelSize, bufferSize);
    }

    // 构造函数，指定名称、通道和缓冲区大小
    public DataCarrier(String name, int channelSize, int bufferSize) {
        this(name, name, channelSize, bufferSize);
    }

    // 构造函数，允许使用环境变量前缀来设置通道和缓冲区大小
    public DataCarrier(String name, String envPrefix, int channelSize, int bufferSize) {
        this(name, envPrefix, channelSize, bufferSize, BufferStrategy.BLOCKING);
    }

    // 构造函数，允许用户自定义缓冲策略，缓冲策略默认为 BLOCKING
    public DataCarrier(String name, String envPrefix, int channelSize, int bufferSize, BufferStrategy strategy) {
        this.name = name;
        // 从环境变量获取缓冲区和通道大小
        bufferSize = EnvUtil.getInt(envPrefix + "_BUFFER_SIZE", bufferSize);
        channelSize = EnvUtil.getInt(envPrefix + "_CHANNEL_SIZE", channelSize);
        // 初始化通道，使用简单的滚动分区策略和指定的缓冲策略
        channels = new Channels<>(channelSize, bufferSize, new SimpleRollingPartitioner<T>(), strategy);
    }

    // 允许自定义缓冲策略的构造函数
    public DataCarrier(int channelSize, int bufferSize, BufferStrategy strategy) {
        this("DEFAULT", "DEFAULT", channelSize, bufferSize, strategy);
    }

    /**
     * 设置新的数据分区策略。它将覆盖当前分区器或默认的分区器。(默认是 {@link SimpleRollingPartitioner})
     *
     * @param dataPartitioner 数据分区器，按照某些规则将数据分配到不同的通道。
     * @return DataCarrier 实例以支持链式调用
     */
    public DataCarrier setPartitioner(IDataPartitioner<T> dataPartitioner) {
        this.channels.setPartitioner(dataPartitioner); // 设置通道的分区策略
        return this;
    }

    /**
     * 向缓冲区生产数据，使用指定的 {@link BufferStrategy}。
     *
     * @return false 表示生产数据失败，数据不会被消费。
     */
    public boolean produce(T data) {
        // 如果驱动存在且未运行，则生产数据失败
        if (driver != null) {
            if (!driver.isRunning(channels)) {
                return false;
            }
        }

        // 保存数据到通道
        return this.channels.save(data);
    }

    /**
     * 设置消费者驱动。消费者将在 {@link DataCarrier#produce} 开始工作时运行。
     *
     * @param consumerClass 消费者类
     * @param num           消费者线程数
     * @param properties    用于初始化消费者的配置。
     */
    public DataCarrier consume(Class<? extends IConsumer<T>> consumerClass,
                               int num,
                               long consumeCycle,
                               Properties properties) {
        // 如果驱动已经存在，关闭之前的驱动
        if (driver != null) {
            driver.close(channels);
        }
        // 创建并启动新的消费者驱动
        driver = new ConsumeDriver<T>(this.name, this.channels, consumerClass, num, consumeCycle, properties);
        driver.begin(channels);
        return this;
    }

    /**
     * 设置消费者驱动。消费者将在 {@link DataCarrier#produce} 开始工作时运行，默认消费周期为 20 毫秒。
     *
     * @param consumerClass 消费者类
     * @param num           消费者线程数
     */
    public DataCarrier consume(Class<? extends IConsumer<T>> consumerClass, int num) {
        return this.consume(consumerClass, num, 20, new Properties());
    }

    /**
     * 设置消费者驱动。消费者将在 {@link DataCarrier#produce} 开始工作时运行。
     *
     * @param consumer 单个消费者实例，所有消费者线程将共享此实例。
     * @param num      消费者线程数
     */
    public DataCarrier consume(IConsumer<T> consumer, int num, long consumeCycle) {
        // 如果之前有驱动，关闭它
        if (driver != null) {
            driver.close(channels);
        }
        // 创建并启动新的消费者驱动
        driver = new ConsumeDriver<T>(this.name, this.channels, consumer, num, consumeCycle);
        driver.begin(channels);
        return this;
    }

    /**
     * 设置消费者驱动。消费者将在 {@link DataCarrier#produce} 开始工作时运行，默认消费周期为 20 毫秒。
     *
     * @param consumer 单个消费者实例，所有消费者线程将共享此实例。
     * @param num      消费者线程数
     */
    public DataCarrier consume(IConsumer<T> consumer, int num) {
        return this.consume(consumer, num, 20);
    }

    /**
     * 为该 DataCarrier 实例设置一个消费者池来管理通道。消费者池可以使用其自己的消费模型调整消费者线程数和吞吐量。
     */
    public DataCarrier consume(ConsumerPool consumerPool, IConsumer<T> consumer) {
        driver = consumerPool; // 使用消费者池作为驱动
        consumerPool.add(this.name, channels, consumer); // 将消费者添加到池中
        driver.begin(channels); // 启动消费者池
        return this;
    }

    /**
     * 关闭所有消费者线程（如果它们正在运行）。注意 {@link BufferStrategy}：如果 {@link BufferStrategy} ==
     * {@link BufferStrategy#BLOCKING}，关闭消费者驱动可能会导致生产阻塞。更好的方式是使用 {@link DataCarrier#consume} 更改消费者驱动。
     */
    public void shutdownConsumers() {
        if (driver != null) {
            driver.close(channels); // 关闭驱动并停止消费者
        }
    }
}
